"""批量清理 YOLO / VOC 标签脚本。

遍历指定根目录下的 `labels`、`obj_train_data` 与 `Annotations` 目录，
删除类别 4 与 5 的标注框，同时将白发类（原 id 6）统一映射为新编号 4。
支持 YOLO `.txt` 与 VOC `.xml` 两种标注格式。
"""

from __future__ import annotations

import argparse
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Set, Tuple


SKIP_YOLO_IDS: Set[int] = {4, 5}
REMAP_YOLO_IDS: Dict[int, int] = {6: 4}
TARGET_YOLO_PARENTS: Set[str] = {"labels", "obj_train_data"}

SKIP_VOC_NAMES: Set[str] = {"4", "5"}
REMAP_VOC_NAMES: Dict[str, str] = {"6": "baifa"}
TARGET_VOC_PARENT: str = "Annotations"


def parse_args(argv: Iterable[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="过滤与重映射检测标注文件")
    parser.add_argument(
        "root",
        type=Path,
        default=Path("VOCdevkit"),
        nargs="?",
        help="需要遍历的根目录，默认指向 VOCdevkit",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅统计将发生的更改，不实际写回文件",
    )
    parser.add_argument(
        "--with-voc",
        action="store_true",
        help="同时处理 VOC Annotations (*.xml) 文件",
    )
    return parser.parse_args(list(argv))


def is_candidate_file(path: Path) -> bool:
    if path.suffix == ".txt" and path.parent.name in TARGET_YOLO_PARENTS:
        return True
    if path.suffix == ".xml" and path.parent.name == TARGET_VOC_PARENT:
        return True
    return False


def process_yolo_file(path: Path) -> Tuple[int, int, List[str]]:
    removed = 0
    remapped = 0
    changed_lines = []

    with path.open("r", encoding="utf-8") as f:
        for raw_line in f:
            line = raw_line.strip()
            if not line:
                changed_lines.append(raw_line)
                continue

            parts = line.split()
            try:
                cls_id = int(parts[0])
            except ValueError:
                # 非数字开头，保留原行
                changed_lines.append(raw_line)
                continue

            if cls_id in SKIP_YOLO_IDS:
                removed += 1
                continue

            new_cls_id = REMAP_YOLO_IDS.get(cls_id, cls_id)
            if new_cls_id != cls_id:
                remapped += 1
            parts[0] = str(new_cls_id)
            newline = "\n" if raw_line.endswith("\n") else ""
            changed_lines.append(" ".join(parts) + newline)

    return removed, remapped, changed_lines


def indent_xml(elem: ET.Element, level: int = 0) -> None:
    indent = "  "
    prefix = "\n" + indent * level
    children = list(elem)
    if children:
        if not elem.text or not elem.text.strip():
            elem.text = prefix + indent
        for child in children:
            indent_xml(child, level + 1)
            if not child.tail or not child.tail.strip():
                child.tail = prefix + indent
    if level and (not elem.tail or not elem.tail.strip()):
        elem.tail = prefix


def process_voc_xml(path: Path) -> Tuple[int, int, Optional[ET.ElementTree]]:
    removed = 0
    remapped = 0

    tree = ET.parse(path)
    root = tree.getroot()

    objects = root.findall("object")
    for obj in list(objects):
        name_elem = obj.find("name")
        if name_elem is None or not name_elem.text:
            continue

        name_text = name_elem.text.strip()

        if name_text in SKIP_VOC_NAMES:
            root.remove(obj)
            removed += 1
            continue

        new_name = REMAP_VOC_NAMES.get(name_text)
        if new_name is not None and new_name != name_text:
            name_elem.text = new_name
            remapped += 1

    if removed == 0 and remapped == 0:
        return 0, 0, None

    indent_xml(root)
    return removed, remapped, tree


def main(argv: Iterable[str]) -> int:
    args = parse_args(argv)
    root: Path = args.root.resolve()

    if not root.exists():
        print(f"根目录不存在: {root}")
        return 1

    total_files = 0
    modified_files = 0
    total_removed = 0
    total_remapped = 0

    for label_file in root.rglob("*"):
        if not is_candidate_file(label_file):
            continue

        if label_file.suffix == ".xml" and not args.with_voc:
            continue

        total_files += 1

        if label_file.suffix == ".txt":
            removed, remapped, new_lines = process_yolo_file(label_file)
            if removed == 0 and remapped == 0:
                continue
            modified_files += 1
            total_removed += removed
            total_remapped += remapped
            if args.dry_run:
                continue
            with label_file.open("w", encoding="utf-8") as fw:
                fw.writelines(new_lines)
        elif label_file.suffix == ".xml":
            removed, remapped, tree = process_voc_xml(label_file)
            if removed == 0 and remapped == 0:
                continue
            modified_files += 1
            total_removed += removed
            total_remapped += remapped
            if args.dry_run or tree is None:
                continue
            tree.write(label_file, encoding="utf-8", xml_declaration=True)

    print("处理完成 OK")
    print(f"扫描标签文件: {total_files}")
    print(f"发生修改文件: {modified_files}")
    print(f"删除的框数: {total_removed}")
    print(f"重映射的框数: {total_remapped}")

    if args.dry_run:
        print("（dry-run 模式，未对文件写入修改）")

    return 0


if __name__ == "__main__":
    # 获取voc_datasets中的所有文件夹
    # 例子:
    # python filter_labels.py F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc --with-voc --dry-run
    sys.exit(main(sys.argv[1:]))
